package com.atguigu.sql

import org.apache.spark.SparkConf
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql.{DataFrame, Encoder, TypedColumn, _}

/**
 * 用户自定义聚合函数（强类型）
 */
object TestSparkSQL_Class_UDAF {

  def main(args: Array[String]): Unit = {

    // 创建配置对象
    val sparkConf = new SparkConf().setAppName("TestSparkSQL").setMaster("local[*]")

    // 创建环境对象
    // enableHiveSupport : 启用hive支持
    val sparkSession = SparkSession
      .builder()
      .config(sparkConf)
      .enableHiveSupport()
      .getOrCreate()

    // 导入隐式转换
    import sparkSession.implicits._

    // 创建聚合函数
    val udaf = new UserAvgAgeClassUDAF

    // 将聚合函数转换为查询列
    val typedColumn: TypedColumn[Student, Double] = udaf.toColumn.name("avgAge")

    // 采用DSL语法实现聚合功能
    val rdd: RDD[(String, Int)] = sparkSession.sparkContext.makeRDD(Array(("zhangsan", 20), ("lisi", 30), ("wangwu", 40)))
    val dataFrame: DataFrame = rdd.toDF("name", "age")

    val dataset: Dataset[Student] = dataFrame.as[Student]

    // 查询指定的列
    dataset.select(typedColumn).show()


    // 关闭资源
    sparkSession.stop()

  }
}

case class Student(name: String, age: Long)

case class AgeBuffer(var sum: Long, var count: Long)

// 声明自定义聚合函数(强类型)
class UserAvgAgeClassUDAF extends Aggregator[Student, AgeBuffer, Double] {

  // 初始化
  override def zero: AgeBuffer = {
    AgeBuffer(0, 0)
  }

  // 更新缓冲区中的数据
  override def reduce(buffer: AgeBuffer, student: Student): AgeBuffer = {
    buffer.sum = buffer.sum + student.age
    buffer.count = buffer.count + 1
    buffer
  }

  // 合并缓冲区
  override def merge(b1: AgeBuffer, b2: AgeBuffer): AgeBuffer = {
    b1.sum = b1.sum + b2.sum
    b1.count = b1.count + b2.count
    b1
  }

  // 计算
  override def finish(buffer: AgeBuffer): Double = {
    buffer.sum.toDouble / buffer.count
  }

  override def bufferEncoder: Encoder[AgeBuffer] = Encoders.product

  override def outputEncoder: Encoder[Double] = Encoders.scalaDouble
}
